#!/usr/bin/python
# -*- coding:utf-8 -*-
# -------------------------------------------------------
# @author  : 龍皓
# @time    : 2025/3/10 21:01
# @function: 数据库连接池工具类
# @version : 1.0.0
# -------------------------------------------------------

import os
import logging
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from core.util.file_util import FileUtil

class DatabasePool:
    def __init__(self, config_path=None, db_type='mysql'):
        self.engine = None
        self.Session = None
        self.db_type = db_type.lower()  # 支持传入 'mysql' 或 'sqlite'
        # 允许通过参数传入配置文件路径，默认为 ../config/settings.yaml
        self.config_path = config_path or os.path.join(os.path.dirname(__file__), '..', 'config', 'settings.yaml')
        self.logger = logging.getLogger(__name__)

    def initialize(self):
        # 读取配置文件
        try:
            config = FileUtil.read_yaml_file(self.config_path)
        except Exception as e:
            self.logger.error(f"读取配置文件失败: {e}")
            raise

        # 获取数据库配置总段
        db_configs = config.get('database')
        if not db_configs:
            raise Exception("配置文件中缺少 'database' 部分")

        # 根据数据库类型分别读取配置参数
        if self.db_type == 'mysql':
            mysql_config = db_configs.get('mysql')
            if not mysql_config:
                raise Exception("配置文件中缺少 MySQL 配置")
            username = mysql_config.get('username')
            password = mysql_config.get('password')
            host = mysql_config.get('host')
            database = mysql_config.get('database')
            connection_string = f'mysql+pymysql://{username}:{password}@{host}/{database}'
        elif self.db_type == 'sqlite':
            sqlite_config = db_configs.get('sqlite')
            if not sqlite_config:
                raise Exception("配置文件中缺少 SQLite 配置")
            database = sqlite_config.get('database')
            connection_string = f'sqlite:///{database}'
        else:
            raise Exception(f"不支持的数据库类型: {self.db_type}")

        # 从配置文件中获取数据库连接池参数
        pool_config = config.get('db_pool', {})
        pool_size = pool_config.get('pool_size', 10)
        max_overflow = pool_config.get('max_overflow', 20)
        pool_timeout = pool_config.get('pool_timeout', 30)
        pool_recycle = pool_config.get('pool_recycle', 1800)
        echo = pool_config.get('echo', False)

        try:
            # 创建引擎并配置连接池
            self.engine = create_engine(
                connection_string,
                pool_size=pool_size,
                max_overflow=max_overflow,
                pool_timeout=pool_timeout,
                pool_recycle=pool_recycle,
                echo=echo
            )
            # 创建会话工厂
            self.Session = sessionmaker(bind=self.engine)
            self.logger.info("数据库连接池初始化成功")
        except Exception as e:
            self.logger.error(f"初始化数据库连接池失败: {e}")
            raise

    def get_session(self):
        if not self.Session:
            raise Exception("Database pool is not initialized.")
        return self.Session()

    def close(self):
        if self.engine:
            self.engine.dispose()
            self.engine = None
            self.Session = None
            self.logger.info("数据库连接池已关闭")

    def __enter__(self):
        if not self.engine:
            self.initialize()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


